package com.btl.service.ai.support.rag.generation.assistant.strategy;

import com.btl.service.ai.support.rag.generation.GenerationModelEnum;
import com.btl.service.ai.support.rag.generation.assistant.AssistantService;
import com.btl.service.ai.support.rag.generation.assistant.GenerationAssistantStrategy;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.IOException;


@Slf4j
@Service
@RequiredArgsConstructor
public class GenerationAssistantOpenAiStrategy extends GenerationAssistantStrategy {

    private final TokenCountEstimator huggingFaceTokenCountEstimator;

    private final ChatModel iOpenAiChatModel;
    private final ChatModel iOpenAiChatFormatResponseModel;
    private final ChatMemoryStore redisChatMemoryStore;
    private final ObjectMapper objectMapper;


    @Override
    public Result<String> generate(Long memoryId, String systemMessage, String userMessage, Query query) {
        AssistantService assistantService = buildAiService(iOpenAiChatModel, redisChatMemoryStore, huggingFaceTokenCountEstimator);
        return this.generate(assistantService, memoryId, systemMessage, userMessage, query);
    }

    @Override
    public <T> Result<T> generateJson(
            Long memoryId,
            String systemMessage,
            String userMessage,
            Query query,
            Class<T> responseType) {
        AssistantService assistantService = buildAiService(
                iOpenAiChatFormatResponseModel,
                redisChatMemoryStore,
                huggingFaceTokenCountEstimator
        );
        Result<String> raw = this.generate(assistantService, memoryId, systemMessage, userMessage, query);
        String cleaned = cleanJson(raw.content());
        if (responseType.equals(String.class)) {
            @SuppressWarnings("unchecked")
            T content = (T) raw.content();
            return copyResult(raw, content);
        }
        T resultObj;
        try {
            resultObj = objectMapper.readValue(cleaned, responseType);
        } catch (IOException e) {
            log.error("JSON parse error: {}", cleaned, e);
            throw new RuntimeException("JSON parse error: " + cleaned, e);
        }
        return copyResult(raw, resultObj);
    }

    @Override
    public <T> Result<T> generateJson(Long memoryId, String systemMessage, String userMessage, Query query, TypeReference<T> typeRef) {
        AssistantService assistantService = buildAiService(
                iOpenAiChatFormatResponseModel,
                redisChatMemoryStore,
                huggingFaceTokenCountEstimator
        );
        Result<String> raw = this.generate(assistantService, memoryId, systemMessage, userMessage, query);
        String cleaned = cleanJson(raw.content());
        T resultObj;
        try {
            resultObj = objectMapper.readValue(cleaned, typeRef);
        } catch (IOException e) {
            log.error("JSON parse error: {}", cleaned, e);
            throw new RuntimeException("JSON parse error: " + cleaned, e);
        }
        // 用 builder 复制原有元数据，只替换 content
        return Result.<T>builder()
                .content(resultObj)
                .tokenUsage(raw.tokenUsage())
                .sources(raw.sources())
                .finishReason(raw.finishReason())
                .toolExecutions(raw.toolExecutions())
                .intermediateResponses(raw.intermediateResponses())
                .finalResponse(raw.finalResponse())
                .build();
    }

    private <T> Result<T> copyResult(Result<String> original, T newContent) {
        return Result.<T>builder()
                .content(newContent)
                .tokenUsage(original.tokenUsage())
                .sources(original.sources())
                .finishReason(original.finishReason())
                .toolExecutions(original.toolExecutions())
                .intermediateResponses(original.intermediateResponses())
                .finalResponse(original.finalResponse())
                .build();
    }

    /**
     * 清理 AI 可能加上的 ```json 包裹
     */
    private String cleanJson(String aiResponse) {
        if (aiResponse == null) return null;
        String cleaned = aiResponse.trim();
        if (cleaned.startsWith("```")) {
            int firstNewline = cleaned.indexOf('\n');
            int lastFence = cleaned.lastIndexOf("```");
            if (firstNewline != -1 && lastFence != -1) {
                cleaned = cleaned.substring(firstNewline + 1, lastFence)
                        .trim();
            }
        }
        return cleaned;
    }


    @Override
    public GenerationModelEnum model() {
        return GenerationModelEnum.OPEN_AI_CHAT_MODEL;
    }
}
